#include "data_utils.h"
#ifndef __CCE_KT_TEST__
#include "acl/acl.h"
extern void multiply_custom_do(uint32_t coreDim, void *l2ctrl, void *stream,
                                uint8_t *x, uint8_t *y, uint8_t *z);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void multiply_custom(GM_ADDR, GM_ADDR, GM_ADDR);
#endif

int32_t main(int32_t argc, char *argv[])
{
    size_t byteSize = 8 * 2048 * sizeof(uint16_t);
    uint32_t blockDim = 8;

#ifdef __CCE_KT_TEST__
    // CPU侧验证
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(byteSize),
            *y = (uint8_t *)AscendC::GmAlloc(byteSize),
            *z = (uint8_t *)AscendC::GmAlloc(byteSize);
    ReadFile("./input/input_x.bin", byteSize, x, byteSize);
    ReadFile("./input/input_y.bin", byteSize, y, byteSize);
    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(multiply_custom, blockDim, x, y, z);
    WriteFile("./output/output_z.bin", z, byteSize);
    AscendC::GmFree(x);
    AscendC::GmFree(y);
    AscendC::GmFree(z);
#else
    // NPU侧验证（简略）
    aclInit(nullptr);
    aclrtSetDevice(0);
    aclrtContext context;
    aclrtCreateContext(&context, 0);
    aclrtStream stream;
    aclrtCreateStream(&stream);

    // 注意：你还需要分配 xDevice, yDevice, zDevice 内存，加载输入数据！
    // 这里只给了核函数调用框架
    multiply_custom_do(blockDim, nullptr, stream, xDevice, yDevice, zDevice);

    aclrtSynchronizeStream(stream);
    // 应拷贝输出并释放内存、销毁Context
    aclFinalize();
#endif
    return 0;
}
